"""ImageUtils类 - 提供图像处理相关的实用功能"""

import os
import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter, ImageDraw, ImageFont
from typing import Tuple, Optional, Union

from rpa_framework.rpa_logging.logger import Logger
from rpa_framework.exceptions.rpa_exceptions import ImageProcessingException, ImageNotFoundException

class ImageUtils:
    """图像工具类，提供图像读取、处理、识别等功能"""
    
    def __init__(self):
        """初始化图像工具类"""
        self.logger = Logger("ImageUtils")
    
    def read_image(self, image_path: str) -> np.ndarray:
        """读取图像文件
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            图像的numpy数组
            
        Raises:
            ImageNotFoundException: 图像文件不存在
            ImageProcessingException: 图像读取失败
        """
        self.logger.debug(f"读取图像: {image_path}")
        
        if not os.path.exists(image_path):
            raise ImageNotFoundException(f"图像文件不存在: {image_path}")
        
        try:
            image = cv2.imread(image_path)
            if image is None:
                raise ImageProcessingException(f"无法读取图像: {image_path}")
            # 转换为RGB格式
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            return image
        except Exception as e:
            self.logger.error(f"读取图像失败: {str(e)}")
            raise ImageProcessingException(f"无法读取图像: {str(e)}") from e
    
    def save_image(self, image: np.ndarray or Image.Image, output_path: str) -> None:
        """保存图像
        
        Args:
            image: 图像数据（numpy数组或PIL Image对象）
            output_path: 输出文件路径
            
        Raises:
            ImageProcessingException: 图像保存失败
        """
        self.logger.debug(f"保存图像: {output_path}")
        
        # 确保目录存在
        directory = os.path.dirname(output_path)
        if directory and not os.path.exists(directory):
            os.makedirs(directory)
        
        try:
            if isinstance(image, np.ndarray):
                # 将numpy数组从RGB转换为BGR（OpenCV格式）
                bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                cv2.imwrite(output_path, bgr_image)
            elif isinstance(image, Image.Image):
                image.save(output_path)
            else:
                raise ValueError("不支持的图像格式")
        except Exception as e:
            self.logger.error(f"保存图像失败: {str(e)}")
            raise ImageProcessingException(f"无法保存图像: {str(e)}") from e
    
    def resize_image(self, image: np.ndarray, width: int, height: int, keep_aspect_ratio: bool = False) -> np.ndarray:
        """调整图像大小
        
        Args:
            image: 原始图像
            width: 目标宽度
            height: 目标高度
            keep_aspect_ratio: 是否保持纵横比
            
        Returns:
            调整后的图像
        """
        if keep_aspect_ratio:
            h, w = image.shape[:2]
            aspect_ratio = w / h
            
            if width / height > aspect_ratio:
                # 以高度为基准
                new_width = int(height * aspect_ratio)
                new_height = height
            else:
                # 以宽度为基准
                new_width = width
                new_height = int(width / aspect_ratio)
            
            resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
            
            # 创建目标大小的空白图像并居中放置
            result = np.full((height, width, 3), 255, dtype=np.uint8)
            y_offset = (height - new_height) // 2
            x_offset = (width - new_width) // 2
            result[y_offset:y_offset+new_height, x_offset:x_offset+new_width] = resized
            return result
        else:
            return cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA)
    
    def crop_image(self, image: np.ndarray, x: int, y: int, width: int, height: int) -> np.ndarray:
        """裁剪图像
        
        Args:
            image: 原始图像
            x: 起始X坐标
            y: 起始Y坐标
            width: 裁剪宽度
            height: 裁剪高度
            
        Returns:
            裁剪后的图像
        """
        h, w = image.shape[:2]
        
        # 确保坐标有效
        x = max(0, x)
        y = max(0, y)
        width = min(width, w - x)
        height = min(height, h - y)
        
        return image[y:y+height, x:x+width]
    
    def convert_to_grayscale(self, image: np.ndarray) -> np.ndarray:
        """转换为灰度图像
        
        Args:
            image: 原始图像
            
        Returns:
            灰度图像
        """
        return cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    
    def enhance_contrast(self, image: np.ndarray, factor: float = 1.5) -> np.ndarray:
        """增强图像对比度
        
        Args:
            image: 原始图像
            factor: 对比度因子（>1增强，<1减弱）
            
        Returns:
            增强后的图像
        """
        pil_image = Image.fromarray(image)
        enhancer = ImageEnhance.Contrast(pil_image)
        enhanced = enhancer.enhance(factor)
        return np.array(enhanced)
    
    def enhance_brightness(self, image: np.ndarray, factor: float = 1.5) -> np.ndarray:
        """增强图像亮度
        
        Args:
            image: 原始图像
            factor: 亮度因子（>1增亮，<1变暗）
            
        Returns:
            增强后的图像
        """
        pil_image = Image.fromarray(image)
        enhancer = ImageEnhance.Brightness(pil_image)
        enhanced = enhancer.enhance(factor)
        return np.array(enhanced)
    
    def find_image_in_image(self, haystack: np.ndarray, needle: np.ndarray, threshold: float = 0.8) -> Optional[Tuple[int, int, int, int]]:
        """在大图中查找小图
        
        Args:
            haystack: 大图
            needle: 小图
            threshold: 匹配阈值
            
        Returns:
            找到的区域 (x, y, width, height)，未找到返回None
        """
        # 转换为灰度图
        haystack_gray = cv2.cvtColor(haystack, cv2.COLOR_RGB2GRAY)
        needle_gray = cv2.cvtColor(needle, cv2.COLOR_RGB2GRAY)
        
        # 使用模板匹配
        result = cv2.matchTemplate(haystack_gray, needle_gray, cv2.TM_CCOEFF_NORMED)
        
        # 查找最佳匹配位置
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        
        self.logger.debug(f"图像匹配最佳相似度: {max_val}")
        
        if max_val >= threshold:
            # 计算匹配区域
            h, w = needle_gray.shape
            top_left = max_loc
            bottom_right = (top_left[0] + w, top_left[1] + h)
            return (top_left[0], top_left[1], w, h)
        
        return None
    
    def add_text_to_image(self, image: np.ndarray, text: str, position: Tuple[int, int], 
                          font_size: int = 20, color: Tuple[int, int, int] = (0, 0, 0),
                          font_path: str = None) -> np.ndarray:
        """在图像上添加文本
        
        Args:
            image: 原始图像
            text: 要添加的文本
            position: 文本位置 (x, y)
            font_size: 字体大小
            color: 文本颜色 (R, G, B)
            font_path: 字体文件路径（可选）
            
        Returns:
            添加文本后的图像
        """
        pil_image = Image.fromarray(image)
        draw = ImageDraw.Draw(pil_image)
        
        # 尝试加载指定字体或使用默认字体
        try:
            if font_path and os.path.exists(font_path):
                font = ImageFont.truetype(font_path, font_size)
            else:
                # 使用默认字体
                font = ImageFont.load_default()
        except Exception:
            font = ImageFont.load_default()
        
        draw.text(position, text, font=font, fill=color)
        return np.array(pil_image)
    
    def highlight_region(self, image: np.ndarray, x: int, y: int, width: int, height: int, 
                        color: Tuple[int, int, int] = (255, 0, 0), thickness: int = 2) -> np.ndarray:
        """高亮显示图像中的区域
        
        Args:
            image: 原始图像
            x: 区域起始X坐标
            y: 区域起始Y坐标
            width: 区域宽度
            height: 区域高度
            color: 边框颜色 (R, G, B)
            thickness: 边框粗细
            
        Returns:
            高亮后的图像
        """
        # 转换为BGR格式用于OpenCV
        bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        
        # 绘制矩形
        cv2.rectangle(bgr_image, (x, y), (x + width, y + height), color[::-1], thickness)
        
        # 转换回RGB格式
        return cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
    
    def compare_images(self, image1: np.ndarray, image2: np.ndarray) -> float:
        """比较两张图像的相似度
        
        Args:
            image1: 第一张图像
            image2: 第二张图像
            
        Returns:
            相似度（0-1）
        """
        # 确保图像大小相同
        if image1.shape != image2.shape:
            image2 = cv2.resize(image2, (image1.shape[1], image1.shape[0]))
        
        # 计算结构相似性指数
        try:
            from skimage.metrics import structural_similarity as ssim
            gray1 = cv2.cvtColor(image1, cv2.COLOR_RGB2GRAY)
            gray2 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
            similarity = ssim(gray1, gray2)
            return similarity
        except ImportError:
            # 如果没有scikit-image，使用简单的像素比较
            diff = np.abs(image1.astype(np.float32) - image2.astype(np.float32)) / 255.0
            mse = np.mean(diff)
            return 1.0 - mse
    
    def apply_blur(self, image: np.ndarray, kernel_size: Tuple[int, int] = (5, 5)) -> np.ndarray:
        """应用模糊效果
        
        Args:
            image: 原始图像
            kernel_size: 模糊内核大小
            
        Returns:
            模糊后的图像
        """
        return cv2.GaussianBlur(image, kernel_size, 0)
    
    def detect_edges(self, image: np.ndarray, low_threshold: int = 50, high_threshold: int = 150) -> np.ndarray:
        """边缘检测
        
        Args:
            image: 原始图像
            low_threshold: 低阈值
            high_threshold: 高阈值
            
        Returns:
            边缘图像
        """
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray, low_threshold, high_threshold)
        # 转换回RGB格式以便一致处理
        return cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
    
    def binarize(self, image: np.ndarray, threshold: int = 128) -> np.ndarray:
        """二值化图像
        
        Args:
            image: 原始图像
            threshold: 阈值
            
        Returns:
            二值化后的图像
        """
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)
        return cv2.cvtColor(binary, cv2.COLOR_GRAY2RGB)
    
    def rotate_image(self, image: np.ndarray, angle: float) -> np.ndarray:
        """旋转图像
        
        Args:
            image: 原始图像
            angle: 旋转角度（度）
            
        Returns:
            旋转后的图像
        """
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        
        # 计算旋转矩阵
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        # 计算旋转后的新尺寸
        cos = np.abs(M[0, 0])
        sin = np.abs(M[0, 1])
        new_w = int((h * sin) + (w * cos))
        new_h = int((h * cos) + (w * sin))
        
        # 调整旋转矩阵以防止裁剪
        M[0, 2] += (new_w / 2) - center[0]
        M[1, 2] += (new_h / 2) - center[1]
        
        # 执行旋转
        rotated = cv2.warpAffine(image, M, (new_w, new_h))
        return rotated
    
    def pil_to_numpy(self, pil_image: Image.Image) -> np.ndarray:
        """将PIL Image转换为numpy数组
        
        Args:
            pil_image: PIL Image对象
            
        Returns:
            numpy数组
        """
        return np.array(pil_image)
    
    def numpy_to_pil(self, numpy_image: np.ndarray) -> Image.Image:
        """将numpy数组转换为PIL Image
        
        Args:
            numpy_image: numpy数组
            
        Returns:
            PIL Image对象
        """
        return Image.fromarray(numpy_image)

if __name__ == "__main__":
    # 示例用法
    image_utils = ImageUtils()
    # 尝试读取测试图像
    try:
        # 创建一个简单的测试图像
        test_image = np.zeros((200, 300, 3), dtype=np.uint8)
        test_image[:, :] = [255, 255, 255]  # 白色背景
        
        # 添加一些内容
        test_image = image_utils.add_text_to_image(test_image, "RPA测试图像", (50, 80), font_size=24, color=(0, 0, 255))
        
        # 保存测试图像
        image_utils.save_image(test_image, "test_image.png")
        print("测试图像已保存为 test_image.png")
        
        # 读取并处理图像
        img = image_utils.read_image("test_image.png")
        gray = image_utils.convert_to_grayscale(img)
        image_utils.save_image(gray, "test_gray.png")
        print("灰度图像已保存为 test_gray.png")
        
        # 增强对比度
        enhanced = image_utils.enhance_contrast(img, 1.5)
        image_utils.save_image(enhanced, "test_enhanced.png")
        print("增强对比度后的图像已保存为 test_enhanced.png")
        
    except Exception as e:
        print(f"示例执行失败: {e}")